# -*- coding: utf-8 -*-            
# @Time : 2023/5/8 16:36
# @FileName: linux_all.py
# @Software: PyCharm
# @Author: zhou chen
import os
import sys
import time

try:
    import paramiko
except ImportError:
    if sys.platform == "win32" or sys.platform == "win64":
        os.system("pip install "
                  + "-i http://pypi.douban.com/simple/ --trusted-host pypi.douban.com "
                  + "paramiko")
    else:
        os.system("pip3 install -i http://pypi.douban.com/simple/ --trusted-host pypi.douban.com "
                  + "paramiko")
    import paramiko

if sys.platform == "win32" or sys.platform == "win64":
    # 获取当前项目的目录
    project_path = os.path.dirname(os.path.dirname(__file__))
    from im_main.handle_md import HandleConfig
else:
    project_path = os.path.abspath("..")
    from handle_md import HandleConfig


def handle_ssh(hostname, port, username, password, config_str):
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    ssh.connect(hostname=hostname, port=port, username=username, password=password)
    ftp = ssh.open_sftp()
    file = ftp.file('/root/IMTest/config/LinuxTestConfig.ini', "w", -1)
    file.write(config_str)

    file.flush()
    file = ftp.file('/root/IMTest/config/LinuxBaseServer.ini', "w", -1)
    with open(project_path + "/config/" + "LinuxBaseServer.ini", "r", encoding="utf-8") as f:
        config_str = f.read()
        file.write(config_str)
    file.flush()

    file = ftp.file('/root/IMTest/config/LinuxServer.ini', "w", -1)
    with open(project_path + "/config/" + "LinuxServer.ini", "r", encoding="utf-8") as f:
        config_str = f.read()
        file.write(config_str)
    file.flush()

    chan = ssh.invoke_shell()
    chan.send('rm -f /root/IMTest/test_file/test_ass.ini \n')
    time.sleep(10)
    chan.send('sh /root/IMTest/im_main/run_client.sh & \n')
    # 检查测试文件
    while True:
        stdin, stdout, stderr = ssh.exec_command("find /root/IMTest/test_file -name 'test_ass.ini'", timeout=10)
        if stdout.read().decode('utf-8'):
            file = ftp.open('/root/IMTest/test_file/test_ass.ini', "r", -1)
            res_list = []
            for link in file:
                res_list = link.split(",")
            ftp.close()
            ssh.close()
            all_send_nub = int(res_list[0])
            all_recv_nub = int(res_list[1])
            all_recv_byte_nub = int(res_list[2])
            avg_send_time = float(res_list[3])
            avg_recv_time = float(res_list[4])
            return all_send_nub, all_recv_nub, all_recv_byte_nub, avg_send_time, avg_recv_time
        else:
            continue


def send_file_to_target(hostname, port, username, password, file_name):
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    ssh.connect(hostname=hostname, port=port, username=username, password=password)
    ftp = ssh.open_sftp()
    file = ftp.file(f'/root/IMTest/test_file/{file_name}', "w", -1)
    with open(project_path + "/test_file/" + "test_ass.ini", "r", encoding="utf-8") as f:
        config_str = f.read()
        file.write(config_str)
    file.flush()
    ftp.close()
    ssh.close()


def write_file(config_str, config_name="LinuxTestConfig.ini"):
    with open(project_path + f"/config/{config_name}", "w", encoding="utf-8") as file:
        file.write(config_str)
        file.close()


def input_test_config(hostnamelist, port, username, password):
    create_test = input("请输入你需要进行的操作 1、创建测试数据；2、IM测试 \n")
    if create_test == "1":
        create_type = input("请输入你需要的创建类型 1、创建用户 2、创建私聊 3、创建群聊 4、创建超级大群 \n")
        if create_type == "1":
            create_user_nub = input("请输入用户数： \n")
            config_str = f"""[TestConfig]
create_test: 1
create_type: 1
create_user_nub: {int(create_user_nub)}
            """
            write_file(config_str)
            return [create_test, create_type, create_user_nub]
        elif create_type == "2":
            create_private_nub = "请输入私聊数量：\n "
            config_str = f"""[TestConfig]
create_test: 1
create_type: 2
create_private_nub: {int(create_private_nub)}
            """
            write_file(config_str)
            return [create_test, create_type, create_private_nub]
        elif create_type == "3":
            group_nub = input("请输入普通群数量： \n")
            group_user_nub = input("请输入普通群 群人数： \n")
            config_str = f"""[TestConfig]
create_test: 1
create_type: 3
group_nub: {int(group_nub)}
group_user_nub: {int(group_user_nub)}
            """
            write_file(config_str)
            return [create_test, create_type, group_nub, group_user_nub]
        elif create_type == "4":
            super_group_nub = input("请输入超大群数量： \n")
            super_group_user_nub = input("请输入超大群人数： \n")
            config_str = f"""[TestConfig]
create_test: 1
create_type: 4
super_group_nub: {int(super_group_nub)}
super_group_user_nub: {int(super_group_user_nub)}
            """
            write_file(config_str)
            return [create_test, create_type, super_group_nub, super_group_user_nub]
        else:
            print("输入内容错误")
            return None
    elif create_test == "2":
        test_type = input("请输入你需要进行的测试类型 1、私聊测试 2、普通群聊测试 3、超级大群测试\n")
        if test_type == "1":
            private_message_nub = int(input("请输入需要测试的会话数(每台设备均会开始输入的会话量)： \n"))
            message_nub = input("请输入每个会话需要发送的消息数量： \n")
            time_sleep = input("请输入休息时间(0 表示不设置休息时间)： \n")
            config_str = f"""[TestConfig]
create_test: 2
test_type: 1
private_message_nub: {int(private_message_nub)}
message_nub: {int(message_nub)}
time_sleep: {float(time_sleep)}
                        """
            write_file(config_str)
            all_send_nubs = 0
            all_recv_nubs = 0
            all_recv_byte_nubs = 0
            avg_send_times = 0.0
            avg_recv_times = 0.0
            i = 0
            for hostname in hostnamelist:
                config_str = f"""[TestConfig]
create_test: 2
test_type: 1
private_message_nub: {int(private_message_nub)}
message_nub: {int(message_nub)}
time_sleep: {float(time_sleep)}
start_nub: {i * private_message_nub}
                """

                all_send_nub, all_recv_nub, all_recv_byte_nub, avg_send_time, avg_recv_time = \
                    handle_ssh(hostname, port, username, password, config_str)
                all_send_nubs += all_send_nub
                all_recv_nubs += all_recv_nub
                all_recv_byte_nubs += all_recv_byte_nub
                avg_send_times += avg_send_time
                avg_recv_times += avg_recv_time
                i += 1
            return [all_send_nubs, all_recv_nubs, all_recv_byte_nubs, avg_send_times, avg_recv_times]
        elif test_type == "2":
            group_nub = input("请输入需要测试的普通群数量： \n")
            send_user_nub = input("请输入每个群发送消息的人数： \n")
            recv_user_nub = input("请输入接收消息的人数： \n")
            send_message_nub = input("请输入需要发送的消息数量： \n")
            time_sleep = input("请输入休息时间(0 表示不设置休息时间)： \n")
            config_str = f"""[TestConfig]
create_test: 2
test_type: 2
group_nub: {int(group_nub)}
send_user_nub: {int(send_user_nub)}
recv_user_nub: {int(recv_user_nub)}
send_message_nub: {int(send_message_nub)}
time_sleep: {float(time_sleep)}
            """
            write_file(config_str)
            all_send_nubs = 0
            all_recv_nubs = 0
            all_recv_byte_nubs = 0
            avg_send_times = 0.0
            avg_recv_times = 0.0
            i = 0
            for hostname in hostnamelist:
                config_str = f"""[TestConfig]
create_test: 2
test_type: 2
group_nub: {int(group_nub)}
send_user_nub: {int(send_user_nub)}
recv_user_nub: {int(recv_user_nub)}
send_message_nub: {int(send_message_nub)}
time_sleep: {float(time_sleep)}
start_nub: {i * (send_user_nub + recv_user_nub)}
                            """
                all_send_nub, all_recv_nub, all_recv_byte_nub, avg_send_time, avg_recv_time = \
                    handle_ssh(hostname, port, username, password, config_str)
                all_send_nubs += all_send_nub
                all_recv_nubs += all_recv_nub
                all_recv_byte_nubs += all_recv_byte_nub
                avg_send_times += avg_send_time
                avg_recv_times += avg_recv_time
                i += 1
            return [all_send_nubs, all_recv_nubs, all_recv_byte_nubs, avg_send_times, avg_recv_times]
        elif test_type == "3":
            super_group_nub = int(input("请输入需要测试的超大群数量： \n"))
            send_user_nub = int(input("请输入发送消息的人数： \n"))
            recv_user_nub = int(input("请输入接收消息的人数：： \n"))
            send_message_nub = int(input("请输入每个群需要发送的消息数量： \n"))
            time_sleep = float(input("请输入休息时间（0表示不设置休息时间： \n"))
            config_str = f"""[TestConfig]
create_test: 2
test_type: 3
super_group_nub: {int(super_group_nub)}
send_user_nub: {int(send_user_nub)}
recv_user_nub: {int(recv_user_nub)}
send_message_nub: {int(send_message_nub)}
time_sleep: {float(time_sleep)}
            """
            write_file(config_str)
            all_send_nubs = 0
            all_recv_nubs = 0
            all_recv_byte_nubs = 0
            avg_send_times = 0.0
            avg_recv_times = 0.0
            i = 0
            for hostname in hostnamelist:
                config_str = f"""[TestConfig]
create_test: 2
test_type: 3
super_group_nub: {int(super_group_nub)}
send_user_nub: {int(send_user_nub)}
recv_user_nub: {int(recv_user_nub)}
send_message_nub: {int(send_message_nub)}
time_sleep: {float(time_sleep)}
start_nub: {i * (int(send_user_nub) + int(recv_user_nub))}
                            """
                all_send_nub, all_recv_nub, all_recv_byte_nub, avg_send_time, avg_recv_time = \
                    handle_ssh(hostname, port, username, password, config_str)
                all_send_nubs += all_send_nub
                all_recv_nubs += all_recv_nub
                all_recv_byte_nubs += all_recv_byte_nub
                avg_send_times += avg_send_time
                avg_recv_times += avg_recv_time
                i += 1
            return [all_send_nubs, all_recv_nubs, all_recv_byte_nubs, avg_send_times, avg_recv_times]
        else:
            print("输入内容错误")
            return None
    else:
        print("输入内容错误")
        return None


# 获取配置信息
def get_server_config():
    if os.path.exists(project_path + f"/config/LinuxBaseServer.ini"):
        print("基础配置存在！")
    else:
        tcp_host = input("请输入IM TCP  地址：\n")
        tcp_port = int(input("请输入IM TCP 端口号：\n"))
        user_url = input("请输入用户url地址：\n")
        shop_url = input("请输入商户url地址：\n")
        merchant_id = input("请输入商户 id：\n")
        app_id = input("请输入商户 app id：\n")
        app_secret = input("请输入商户 app 密钥：\n")
        config_str = f"""[server]
tcp_host: {tcp_host}
tcp_port: {tcp_port}
user_url: {user_url}
shop_url: {shop_url}
merchant_id: {merchant_id}
app_id: {app_id}
app_secret: {app_secret}
            """
        write_file(config_str, "LinuxBaseServer.ini")
    if os.path.exists(project_path + f"/config/LinuxServer.ini"):
        reset = input("远程设备信息已存在是否修改：（Y/N）")
        if reset.upper() == "N":
            handle_config = HandleConfig("LinuxServer.ini")
            host_name_list = handle_config.get_value("linux_server", "host_name_list")
            host_port = int(handle_config.get_value("linux_server", "host_port"))
            user_name = handle_config.get_value("linux_server", "user_name")
            pass_word = handle_config.get_value("linux_server", "pass_word")
            return eval(host_name_list), host_port, user_name, pass_word

    host_name_list_str = input("请输入测试的IP， 多个 IP 请以逗号进行区分：\n")
    print(f"你输入的内容为： {host_name_list_str}")
    host_name_list = host_name_list_str.split(",")
    host_port = int(input("请输入端口号:\n"))
    user_name = input("请输入设备的用户名：\n")
    pass_word = input("请输入设备密码：\n")
    print(f"配置内容为： IP： {host_name_list}, PORT : {host_port}, user name : {user_name}, password: {pass_word}\n")
    save_config = input("是否需要保存以上配置到 本地以便下次使用： （Y/N）\n")
    if save_config.upper() == "Y":
        config_str = f"""[linux_server]
host_name_list: {host_name_list}
host_port: {host_port}
user_name: {user_name}
pass_word: {pass_word}
                """
        print(config_str)
        write_file(config_str, "LinuxServer.ini")
    return host_name_list, host_port, user_name, pass_word


if __name__ == "__main__":
    try:
        host_name_lists, host_ports, user_names, pass_words = get_server_config()
        input_res = input_test_config(host_name_lists, host_ports, user_names, pass_words)
        print(input_res)
    except Exception as err:
        print(err)
